package com.koicarp.agent.rag;

import java.util.ArrayList;
import java.util.List;

import cn.hutool.json.JSONArray;
import cn.hutool.json.JSONObject;
import cn.hutool.json.JSONUtil;
import dev.langchain4j.data.document.Document;
import dev.langchain4j.data.document.DocumentSplitter;
import dev.langchain4j.data.document.Metadata;
import dev.langchain4j.data.segment.TextSegment;

/**
 * 自定义隐患切分类
 */
public class TroubleJsonDocumentSplitter implements DocumentSplitter {

	@Override
	public List<TextSegment> split(Document document) {
		String text = document.text();
		JSONArray array = JSONUtil.parseArray(text);
		List<TextSegment> textSegments = new ArrayList<>();
		array.forEach(a->{
			String jsonStr = JSONUtil.toJsonStr(a);
			Metadata metadata = new Metadata();
			JSONObject object = JSONUtil.parseObj(a);
			metadata.put("创建时间", object.get("创建时间").toString());
			TextSegment textSegment = new TextSegment(jsonStr, metadata);
			textSegments.add(textSegment);
		});
		return textSegments;
	}

}
